package cn.cuilan.other.tree;

/**
 * 红黑树实现
 * <p>
 * 1.每个节点必须是红色或黑色；
 * 2.根节点必须是黑色；
 * 3.每个叶子节点都是黑色的空节点；
 * 4.如果节点是红色，则它的子节点必须是黑色（反之不一定）；红红不能相连
 * 5.从根节点到叶节点或空子节点的每条路径，必须包含相同数目的黑色节点（即：相同的黑色高度）。
 *
 * @author zhang.yan
 * @since 2021/3/19
 */
public class RedBlackTree<T extends Comparable<T>> {

    /**
     * 根节点
     */
    private volatile Node<T> root;

    // TODO delete
    public Node<T> getRoot() {
        return root;
    }

    /**
     * 插入
     *
     * @param obj 需要实现Comparable接口
     */
    public void insert(T obj) {
        if (root == null) {
            // 根节点必须是黑色
            root = new Node<>(null, null, null, false, obj);
            return;
        }

        Node<T> p = root;
        while (true) {
            int compare = p.obj.compareTo(obj);
            // 0-相等，不做任何操作，返回
            if (compare == 0) {
                return;
            }

            // 1-大于，放入左节点
            if (compare > 0) {
                if (p.left == null) {
                    Node<T> x = new Node<>(p, null, null, false, obj);
                    p.left = x;
                    balanceInsertion(x);
                    return;
                } else {
                    p = p.left;
                }
            }
            // -1-小于，放入右节点
            else {
                if (p.right == null) {
                    Node<T> x = new Node<>(p, null, null, false, obj);
                    p.right = x;
                    balanceInsertion(x);
                    return;
                } else {
                    p = p.right;
                }
            }
        }
    }

    /**
     * 插入后自平衡
     *
     * @param x 插入的节点
     * @return
     */
    private Node<T> balanceInsertion(Node<T> x) {
        x.red = true;
        for (Node<T> xp, xpp, xppl, xppr; ; ) {

            // 如果父节点为空，说明此节点为根节点
            xp = x.parent;
            if (xp == null) {
                // 变黑色，返回
                x.red = false;
                return x;
            }

            // 如果父节点是黑色，或祖父节点为空，说明父节点为根节点
            xpp = xp.parent;
            if (!xp.red || xpp == null) {
                return root;
            }

            // 如果父节点是祖父节点的左节点
            if (xp == (xppl = xpp.left)) {
                // 如果祖父节点的右节点不为空，且为红色
                if ((xppr = xpp.right) != null && xppr.red) {
                    // 变黑色
                    xppr.red = false;
                    xp.red = false;
                    xpp.red = true;
                    x = xpp;
                } else {
                    if (x == xp.right) {
                        // 左旋
                        root = rotateLeft(x = xp);
                        xpp = (xp = x.parent) == null ? null : xp.parent;
                    }
                    if (xp != null) {
                        xp.red = false;
                        if (xpp != null) {
                            xpp.red = true;
                            root = rotateRight(xpp);
                        }
                    }
                }
            }
            // 如果父节点是祖父节点的右节点
            else {
                // 如果祖父节点的左节点不为空，且为红色
                if (xppl != null && xppl.red) {
                    // 变黑色
                    xppl.red = false;
                    xp.red = false;
                    xpp.red = true;
                    x = xpp;
                } else {
                    if (x == xp.left) {
                        // 右旋
                        root = rotateRight(x = xp);
                        xpp = (xp = x.parent) == null ? null : xp.parent;
                    }
                    if (xp != null) {
                        xp.red = false;
                        if (xpp != null) {
                            xpp.red = true;
                            root = rotateLeft(xpp);
                        }
                    }
                }
            }
        }
    }

    /**
     * 左旋
     *
     * @param p 当前节点
     * @return 返回左旋调整后的root节点
     */
    private Node<T> rotateLeft(Node<T> p) {
        Node<T> r, pp, rl;
        // p节点不能为空，且p的右节点不为空，并将p的右节点赋值给r
        if (p != null && (r = p.right) != null) {

            // 判断 r 的左节点是否为空
            if ((rl = p.right = r.left) != null) {
                rl.parent = p;
            }

            // 是否为根节点，是否需要变色
            if ((pp = r.parent = p.parent) == null) {
                (root = r).red = false;
            }
            // 判断 p 是 pp 的左节点还是右节点
            else if (pp.left == p) {
                // 左节点
                pp.left = r;
            } else {
                // 右节点
                pp.right = r;
            }

            // 将 p 变更为 r 的左节点
            r.left = p;
            p.parent = r;
        }
        return root;
    }

    /**
     * 右旋
     *
     * @param p 当前节点
     * @return 返回右旋调整后的root节点
     */
    private Node<T> rotateRight(Node<T> p) {
        Node<T> l, pp, lr;
        // p 节点不能为空，且 p 的左节点不为空，并将 p 的左节点赋值给 l。
        if (p != null && (l = p.left) != null) {

            // 判断 l 的右节点是否为空
            if ((lr = p.left = l.right) != null) {
                lr.parent = p;
            }

            // 是否为根节点，是否需要变色
            if ((pp = l.parent = p.parent) == null) {
                (root = l).red = false;
            }
            // 判断 p 是 pp 的左节点还是右节点
            else if (pp.right == p) {
                pp.right = l;
            } else {
                pp.left = l;
            }

            // 将 p 变更为 l 的右节点
            l.right = p;
            p.parent = l;
        }
        return root;
    }

    /* ------------------------------------------------------------ */

    /**
     * Node节点
     *
     * @param <T> 可比较的类型
     */
    private static class Node<T> {
        // 连接父节点
        Node<T> parent;
        // 连接左节点
        Node<T> left;
        // 连接右节点
        Node<T> right;
        // true:红,false:黑
        boolean red;
        T obj;

        public Node(Node<T> parent, Node<T> left, Node<T> right, boolean red, T obj) {
            this.parent = parent;
            this.left = left;
            this.right = right;
            this.red = red;
            this.obj = obj;
        }

//        public Node<T> getParent() {
//            return parent;
//        }

        public Node<T> getLeft() {
            return left;
        }

        public Node<T> getRight() {
            return right;
        }

//        public boolean isRed() {
//            return red;
//        }

        public T getObj() {
            return obj;
        }
    }
}
